from openai import OpenAI
import os
from datasets import load_dataset

def chatNow(model="gpt-3.5-turbo-instruct",mode='balance',prompt=None):
    #根据mode来调整temperature及presence_penalty

    if mode == 'balance':
        temperature = 1
        presence_penalty = 0
    elif mode == 'precision':
        temperature = 0.8
        presence_penalty = 2
    elif mode == 'creativity':
        temperature = 1.2
        presence_penalty = -1

    def chat(prompt):
        try:
            client = OpenAI(api_key="sk-yrpeg76LZt53LhyTcue8ojFcvRkzs53fapexPklNqR5KPKae", base_url="https://api.fe8.cn/v1")
            response = client.completions.create(model=model, prompt=prompt, max_tokens=1000, temperature=temperature, presence_penalty=presence_penalty)

            return response.choices[0].text.strip()
        except Exception as e:
            return "broken"

    print(chat(prompt))

#chatNow(mode="balance",model= "gpt-3.5-turbo-instruct",prompt="""Q:“艾米需要4分钟才能爬到滑梯顶部，她花了1分钟才滑下来，水滑梯将在15分钟后关闭，请问在关闭之前她能滑多少次？"
#A：为了解决“在关闭之前她能滑多少次？”这个问题，我们首先要解决的问题是""")

def downloadDataSet():
    dataset = load_dataset("scan","simple",trust_remote_code=True)
    print(dataset)

# downloadDataSet()

Command1 = 'look thrice after jump'
Action1 = 'JUMP LOOK LOOK LOOK'
Command2 = 'run left and walk'
Action2 = 'TURN LEFT RUN WALK'
Command3 = 'look opposite right'
Action3 = 'TURN RIGHT TURN RIGHT LOOK'

# #.zero-shot
# chatNow(prompt=Command1)
# #输出
# """
# <|diff_marker|> --- lib/util/jump.js
# -    if (document.selection) range = document.selection.createRange().dataObject;
# <|diff_marker|> 1012
# +    if (document.selection) range = document.selection.createRange();
# """

# # few-shot
# prompt = f"""
# Q: {Command1}
# A: {Action1}
# Q: {Command2}
# A: {Action2}
# Q: {Command3}
# A:
# """
# chatNow(prompt=prompt)

# zero-shot-ltm
prompt = f"""
In order to translate ‘{Command1}’,we need to first solve
"""
chatNow(prompt=prompt)
